import pandas as pd
import json
import os
import matplotlib.pyplot as plt
from tqdm import tqdm

# 定义路径
preprocessed_dir = 'fujian/fujian1/preprocessed'
original_data_dir = 'fujian/fujian1/json_output'
charts_dir = 'fujian/fujian1/charts_sample_10groups'

# 创建保存图表的文件夹（如果不存在的话）
os.makedirs(charts_dir, exist_ok=True)

# 找到所有处理后的 JSON 文件
processed_files = [f for f in os.listdir(preprocessed_dir) if f.endswith('.json')]

# 创建一个字典来存储均值
qty_means = {}

# tqdm 进度条，用于遍历处理后的文件，计算均值
for processed_file in tqdm(processed_files, desc="Calculating averages"):
    with open(os.path.join(preprocessed_dir, processed_file), 'r') as f:
        records = json.load(f)
        
        # 计算 qty 的均值
        avg_qty = sum(record['qty_smoothed'] for record in records) / len(records)
        qty_means[processed_file] = avg_qty

# 找出均值前 10 大的文件
top_10_files = sorted(qty_means, key=qty_means.get, reverse=True)[:10]

# tqdm 进度条，用于为每个文件绘制图表
for processed_file in tqdm(top_10_files, desc="Generating charts"):
    # 读取处理后的数据
    with open(os.path.join(preprocessed_dir, processed_file), 'r') as f:
        processed_records = json.load(f)

    # 获取相应的原始数据文件名
    original_file = processed_file.replace('processed_', '').replace('.json', '.json')
    
    # 读取原始数据
    with open(os.path.join(original_data_dir, original_file), 'r') as f:
        original_records = json.load(f)

    # 提取日期和 qty 数据
    processed_dates = [record['date'] for record in processed_records]
    processed_qty = [record['qty_smoothed'] for record in processed_records]
    original_qty = [record['qty'] for record in original_records]

    # 绘制图表
    plt.figure(figsize=(12, 6))
    plt.plot(processed_dates, processed_qty, label='Processed (Smoothed)', color='blue')
    plt.plot(processed_dates, original_qty, label='Original', color='red', linestyle='--')
    
    # 图表设置
    plt.title(f"Comparison of Processed and Original qty - {original_file}")
    plt.xlabel("Date")
    plt.ylabel("Quantity")
    plt.xticks(rotation=45)
    plt.legend()
    plt.tight_layout()

    # 保存图表
    plt.savefig(f"{charts_dir}/{original_file.replace('.json', '.png')}")
    plt.close()

print("All charts saved in 'charts_sample_10groups' directory.")
